import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import datetime
import math

from sklearn.linear_model import LinearRegression
from sklearn import preprocessing
from sklearn.model_selection import train_test_split
import tushare as ts
 
data =ts.get_hist_data('600519')
data.to_csv('./600519.csv')
data=pd.read_csv('./600519.csv')
df = data[['open','high','low','volume','close']]

featureDatas = df[['open','high','low','volume']]
feature = featureDatas.values
target = np.array(df['close'])
 

feature_train, feature_test, target_train, target_test = train_test_split(feature,target,test_size = 0.05)
lrtoot = LinearRegression() 
lrtoot.fit(feature_train,target_train)  
 
predictByTest = lrtoot.predict(feature_test)
predictDays = int(math.ceil(0.05 * len(df)))  
 

index = 0
while index < len(data) - predictDays:
    df.loc[index,'predictValue'] = data.loc[index,'close']    
    df.loc[index,'date'] = data.loc[index,'date']             
    index = index + 1
    

predictedCnt = 0
while predictedCnt < predictDays:
    df.loc[index,'predictValue'] = predictByTest[predictedCnt]   
    df.loc[index,'date'] = data.loc[index,'date']     
    
    predictedCnt =predictedCnt + 1
    index =index + 1

plt.figure(figsize=(30,10))
 
df['predictValue'].plot(color='red',label='predict data',fontsize=30)
df['close'].plot(color='blue',label='real data',fontsize=30)
 
 
plt.legend(loc = 'best',fontsize=40)
 

major_index = df.index[df.index%30==0]
major_xtics = df['date'][df.index%30==0]
plt.xticks(major_index,major_xtics)
plt.setp(plt.gca().get_xticklabels(), rotation=30)
 

plt.grid(linestyle='-.')
 
plt.show()